package parser;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import edu.stanford.nlp.ling.CoreAnnotations.LemmaAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.NamedEntityTagAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.PartOfSpeechAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.SentencesAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.TextAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.TokensAnnotation;
import edu.stanford.nlp.ling.CoreLabel;
import edu.stanford.nlp.pipeline.Annotation;
import edu.stanford.nlp.util.CoreMap;

public class StanfordAnnotation extends ParserAnnotation {

	private StanfordPipeline p;
	private Annotation ann;

	private List<String> sents;
	private List<List<String>> lemmas;
	private List<List<String>> pos;
	private List<List<String>> tokens;
	private List<List<String>> ner;
	private List<List<Integer>> indexes;

	StanfordAnnotation(StanfordPipeline p, String document) {
		super(document);
		this.p = p;
		lemmas = new ArrayList<List<String>>();
		pos = new ArrayList<List<String>>();
		tokens = new ArrayList<List<String>>();
		ner = new ArrayList<List<String>>();
		indexes = new ArrayList<List<Integer>>();
	}

	@Override
	protected void parseImplementation() {
		ann = p.process(document);
		sents = makeSentences();
		for (int i = 0; i < sents.size(); i++) {
			lemmas.add(makeLemmas(i));
			pos.add(makePOS(i));
			tokens.add(makeTokens(i));
			ner.add(makeNER(i));
			indexes.add(makeIndexes(i));
		}
		for (int i = 0; i < sents.size(); i++) {
			localizeIndexes(indexes.get(i));
		}
	}

	private void localizeIndexes(List<Integer> index) {
		int start = index.get(0);
		for (int i = 0; i < index.size(); i++)
			index.set(i, index.get(i) - start);
	}

	private List<String> makeSentences() {
		if (skipSentDetection)
			return Arrays.asList(document);
		List<CoreMap> cms = ann.get(SentencesAnnotation.class);
		List<String> sentences = new ArrayList<String>(cms.size());
		for (int i = 0; i < cms.size(); i++) {
			String sent = cms.get(i).toString();
			// sent = sent.replaceAll("[\n\r\t]", " ");
			// sent = sent.replaceAll(" +", " ");
			sentences.add(sent);
		}
		return sentences;
	}

	public int getNbrSentences() {
		if (skipSentDetection)
			return 1;
		return sents.size();
	}

	private List<CoreLabel> getLabels(int sent) {
		if (skipSentDetection) {
			List<CoreMap> cms = ann.get(SentencesAnnotation.class);
			List<CoreLabel> cls = new ArrayList<CoreLabel>();
			for (int i = 0; i < cms.size(); i++)
				cls.addAll(cms.get(i).get(TokensAnnotation.class));
			return cls;
		}
		CoreMap map = ann.get(SentencesAnnotation.class).get(sent);
		List<CoreLabel> cls = map.get(TokensAnnotation.class);
		return cls;
	}

	private List<Integer> makeIndexes(int sent) {
		List<CoreLabel> cls = getLabels(sent);
		Integer[] tokens = new Integer[cls.size() * 2];
		for (int i = 0; i < cls.size(); i++) {
			tokens[i * 2] = cls.get(i).beginPosition();
			tokens[i * 2 + 1] = cls.get(i).endPosition();
		}
		return Arrays.asList(tokens);
	}

	private List<String> makeTokens(int sent) {
		List<CoreLabel> cls = getLabels(sent);
		String[] tokens = new String[cls.size()];
		for (int i = 0; i < cls.size(); i++) {
			tokens[i] = cls.get(i).getString(TextAnnotation.class);
			tokens[i] = reversePTB(tokens[i]);
		}
		return Arrays.asList(tokens);
	}

	private List<String> makePOS(int sent) {
		List<CoreLabel> cls = getLabels(sent);
		String[] pos = new String[cls.size()];
		for (int i = 0; i < cls.size(); i++)
			pos[i] = cls.get(i).get(PartOfSpeechAnnotation.class);
		return Arrays.asList(pos);
	}

	private List<String> makeLemmas(int sent) {
		List<CoreLabel> cls = getLabels(sent);
		String[] lemmas = new String[cls.size()];
		for (int i = 0; i < cls.size(); i++) {
			lemmas[i] = cls.get(i).get(LemmaAnnotation.class);
			if (lemmas[i] == null || lemmas[i].isEmpty())
				lemmas[i] = "-unknown-";
			lemmas[i] = reversePTB(lemmas[i]);
			lemmas[i] = normalize(lemmas[i]);
		}
		return Arrays.asList(lemmas);
	}

	private List<String> makeNER(int sent) {
		List<CoreLabel> cls = getLabels(sent);
		List<String> ner = new ArrayList<String>();
		for (int i = 0; i < cls.size(); i++) {
			String nertag = cls.get(i)
					.getString(NamedEntityTagAnnotation.class);
			if (nertag.contains("-NBR-")) {
				String token = reversePTB(cls.get(i).getString(
						TextAnnotation.class));
				Pattern nbr = Pattern.compile(".*?(1?[0-9])(?:\\.[0-9.]+)*");
				Matcher mat = nbr.matcher(token);
				if (mat.find())
					token = mat.group(1);
				else
					token = "";
				ner.add(token + "/" + nertag.split("/")[1]);
			} else
				ner.add(nertag);
		}
		return ner;
	}

	private String reversePTB(String str) {
		String source = str.replace("-LRB-", ")");
		source = source.replace("-lrb-", ")");
		source = source.replace("-RRB-", ")");
		source = source.replace("-rrb-", ")");
		source = source.replace("''", "\"");
		source = source.replace("´´", "´´");
		source = source.replace("``", "``");
		source = source.replace('`', '\'');
		source = source.replace('´', '`');
		source = source.replace("\\/", "/");
		return source;
	}

	@Override
	public List<String> getSentences() {
		return sents;
	}

	@Override
	public List<String> getTokens(int sent) {
		return tokens.get(sent);
	}

	@Override
	public List<String> getPOS(int sent) {
		return pos.get(sent);
	}

	@Override
	public List<String> getLemmas(int sent) {
		return lemmas.get(sent);
	}

	public List<String> getNER(int sent) {
		return ner.get(sent);
	}

	@Override
	public List<Integer> getTokenIndexes(int sent) {
		return indexes.get(sent);
	}
}
